import random
import numpy as np
import time
import pickle
import math
random.seed(int(time.time()))
# -----------------------------
# 1. 参数区（可自行调整）
# -----------------------------
num_floors = 3          # 楼层数
floor_width = 30       # 每层宽（格子数）
floor_length = 50      # 每层长（格子数）
floor_height = 4      # 楼梯间高度 (m)
people_per_floor = 100   # 每层初始人数
stairs_per_floor = 10 # 每层上楼梯数量
exits_on_ground = 10     # 一楼出口数量
cell_size = 1.0         # 每格面积(m²)
obstacle_size = (2,10)   # 障碍物尺寸
obstacle_attempt = 10    # 放障碍物尝试 (2*4) 或 (4*2) 形状
speed_max = 1.2         # 平路运动最高速 (m/s)
speed_min = 0.8         # 平路运动最低速 (m/s)
speed_floor_down_radio = 1.0 #下楼速度和平路运动的比值
time_update = 1     # 每帧1秒 (每1秒更新一次)

# -----------------------------
# 2. 初始化楼层与楼梯
# -----------------------------
floors = np.zeros((num_floors, floor_width, floor_length), dtype=int)
stairs_up = {}
stairs_down = {}

for f in range(num_floors):
    floors[f,0,:] = 1; floors[f,-1,:] = 1
    floors[f,:,0] = 1; floors[f,:,-1] = 1
    stairs_up[f] = []
    stairs_down[f] = []

# 随机布置上楼梯和下楼梯（不同位置）
def place_stairs(num_up):
    up_list = []
    occupied = set()
    for _ in range(num_up):
        while True:
            i,j = random.randint(1,floor_width-2), random.randint(1,floor_length-2)
            if (i,j) not in occupied:
                up_list.append((i,j))
                occupied.add((i,j))
                break
    return up_list

downs = None
for f in range(num_floors):
    if f != num_floors - 1:
        ups = place_stairs(stairs_per_floor)
        stairs_up[f] = ups
        for i,j in ups:
            floors[f,i,j] = 4  # 4代表上楼梯        
    if f != 0:
        stairs_down[f] = downs
        for i,j in downs:
            floors[f,i,j] = 5  # 5代表下楼梯        
    downs = ups

# 随机放障碍物
for f in range(num_floors):
    for i in range(obstacle_attempt):
        if random.random() <= 0.5:
            h, w = obstacle_size
        else:
            w, h = obstacle_size
        rand_i = random.randint(0,floor_width-h)
        rand_j = random.randint(0,floor_length-w)            
        if sum(sum(floors[f,rand_i:rand_i+h,rand_j:rand_j+w])) == 0:
            floors[f,rand_i:rand_i+h,rand_j:rand_j+w] = 1

# 一楼随机放出口，不能和楼梯重合
exits = []
occupied_exit = set(stairs_up[0])
for _ in range(exits_on_ground):
    while True:
        side = random.choice(['top','bottom','left','right'])
        if side == 'top':
            pos = (0, random.randint(1,floor_length-2))
        elif side == 'bottom':
            pos = (floor_width-1, random.randint(1,floor_length-2))
        elif side == 'left':
            pos = (random.randint(1,floor_width-2), 0)
        else:
            pos = (random.randint(1,floor_width-2), floor_length-1)
        if pos not in occupied_exit:
            floors[0,pos[0],pos[1]] = 3  # 3代表出口
            exits.append((0,pos))
            break

# -----------------------------
# 3. 初始化人员
# -----------------------------
people = []
for f in range(num_floors):
    count = 0
    while count < people_per_floor:
        i,j = random.randint(1,floor_width-2), random.randint(1,floor_length-2)
        if floors[f,i,j] in [0,4,5]:  # 空地或楼梯可以初始放人
            # people.append([f,(i,j),random.uniform(speed_min,speed_max), math.inf])
            people.append([f,(i,j),math.inf])
            count +=1

data_floors = (\
    num_floors, \
    floor_width, \
    floor_length, \
    floor_height, \
    exits_on_ground, \
    cell_size, \
    speed_floor_down_radio, \
    time_update, \
    floors, \
    stairs_up, \
    stairs_down, \
    exits, \
    people)

with open("data/data_floors_0427.pkl", "wb") as haddle:
    pickle.dump(data_floors, haddle)